from torch import nn
from torch.autograd import Function
from torch.autograd.function import once_differentiable

from taa_core._C import star_dcn_offset


class StarDCNOffsetF(Function):
    @staticmethod
    def forward(ctx, bbox_pred, dcn_base_offset, grad_mul_idx, signedK, num_dconv_points, gradient_mul, stride):
        return star_dcn_offset(bbox_pred, dcn_base_offset, grad_mul_idx, signedK, num_dconv_points, gradient_mul, stride)

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_outputs):
        return None


star_dcn_offset_fun = StarDCNOffsetF.apply


class StarDcnOffset(nn.Module):
    def __init__(self):
        super(StarDcnOffset, self).__init__()

    def forward(self, bbox_pred, dcn_base_offset, grad_mul_idx, signedK, num_dconv_points, gradient_mul, stride):
        return star_dcn_offset_fun(bbox_pred, dcn_base_offset, grad_mul_idx, signedK, num_dconv_points, gradient_mul, stride)